package com.ctjsoft.eureka;

import com.netflix.appinfo.ApplicationInfoManager;
import com.netflix.appinfo.InstanceInfo;
import com.netflix.discovery.EurekaClient;
import com.netflix.discovery.EurekaClientConfig;
import com.netflix.discovery.shared.Application;
import com.netflix.eureka.EurekaServerConfig;
import com.netflix.eureka.lease.Lease;
import com.netflix.eureka.registry.PeerAwareInstanceRegistryImpl;
import com.netflix.eureka.resources.ServerCodecs;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.cloud.netflix.eureka.server.event.EurekaInstanceCanceledEvent;
import org.springframework.cloud.netflix.eureka.server.event.EurekaInstanceRegisteredEvent;
import org.springframework.cloud.netflix.eureka.server.event.EurekaInstanceRenewedEvent;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ApplicationEvent;
import org.springframework.util.CollectionUtils;

import java.util.List;

/**
 * 设置 Eureka 注册中心 IP 白名单，只有白名单内的 IP 允许注册。
 * 亲测在 Spring boot 2.1.3 以及 Spring boot 1.5.2 版本下没有问题.
 * <p>
 * 1、Eureka 注册中心在实际使用中，经常发生陌生 IP 注册上来，导致服务无法正常访问。
 * 2、在一般的强制服务实例下线/关闭的方法中，通常都是需要注册方改动代码进行支持。而显然这种方式
 * 是比较被动的，因为有时候都不知道是谁注册上来的。
 * 3、所以最直接的方式是直接修改 Eureka 注册中心服务端，只允许指定 IP 进行访问。
 * <p>
 * PeerAwareInstanceRegistryImpl：Eureka 实例注册表实现，用于管理服务节点注册.
 *
 * @author wangMaoXiong
 * @version 1.0
 * @date 2022/8/26 14:19
 */
public class EurekaInstanceRegistry extends PeerAwareInstanceRegistryImpl implements ApplicationContextAware {

    private ApplicationContext applicationContext;

    private int defaultOpenForTrafficCount;

    /**
     * 允许注册到 Eureka 的 IP 白名单列表。
     * 如果为空，则表示不进行限制，只有不为空时，才只允许指定 IP 进行注册.
     */
    private List<String> allowedRegisteredIpAddress;

    private static final Logger log = LoggerFactory.getLogger(EurekaInstanceRegistry.class);

    /**
     * 构造器
     *
     * @param serverConfig
     * @param clientConfig
     * @param serverCodecs
     * @param eurekaClient
     * @param defaultOpenForTrafficCount
     * @param allowedRegisteredIpAddress
     */
    public EurekaInstanceRegistry(EurekaServerConfig serverConfig,
                                  EurekaClientConfig clientConfig,
                                  ServerCodecs serverCodecs,
                                  EurekaClient eurekaClient,
                                  int defaultOpenForTrafficCount,
                                  List<String> allowedRegisteredIpAddress) {
        super(serverConfig, clientConfig, serverCodecs, eurekaClient);
        this.defaultOpenForTrafficCount = defaultOpenForTrafficCount;
        this.allowedRegisteredIpAddress = allowedRegisteredIpAddress;
    }

    @Override
    public void setApplicationContext(ApplicationContext context) throws BeansException {
        this.applicationContext = context;
    }


    @Override
    public void openForTraffic(ApplicationInfoManager applicationInfoManager, int count) {
        super.openForTraffic(applicationInfoManager, count == 0 ? this.defaultOpenForTrafficCount : count);
    }

    /**
     * 用给定的持续时间注册新的实例。
     *
     * @param instanceInfo  ：实例信息
     * @param leaseDuration ：租赁期限
     * @param isReplication ：是否复制
     */
    @Override
    public void register(InstanceInfo instanceInfo, int leaseDuration, boolean isReplication) {
        log.info("\n Request Register: \n\tid={}, \n\thostName={}, \n\tappName={}, \n\tvipAddress={}, \n\tipAddr={}, \n\tport={}, \n\tleaseDuration={}, \n\tisReplication={}.",
                instanceInfo.getId(), instanceInfo.getHostName(), instanceInfo.getAppName(), instanceInfo.getVIPAddress(),
                instanceInfo.getIPAddr(), instanceInfo.getPort(), leaseDuration, isReplication);

        publishEvent(new EurekaInstanceRegisteredEvent(this, instanceInfo, leaseDuration, isReplication));

        // 如果IP白名单为空，则表示不进行限制
        // IP白名单不为空时，只允许指定 IP 进行注册.
        if (CollectionUtils.isEmpty(allowedRegisteredIpAddress) || allowedRegisteredIpAddress.contains(instanceInfo.getIPAddr())) {
            log.info("\n IP Address Allows Registration: {}", instanceInfo.getIPAddr());
            super.register(instanceInfo, leaseDuration, isReplication);
            return;
        }
        log.warn("\n IP Address Not Allows Registration: {}", instanceInfo.getIPAddr());
    }

    /**
     * 实例注册
     *
     * @param instanceInfo  ：实例信息
     * @param isReplication ：是否复制
     */
    @Override
    public void register(final InstanceInfo instanceInfo, final boolean isReplication) {
        log.info("\n Request Register: \n\tid={}, h\n\tostName={}, \n\tappName={}, \n\tvipAddress={}, \n\tipAddr={}, \n\tport={}, \n\tisReplication={}.",
                instanceInfo.getId(), instanceInfo.getHostName(), instanceInfo.getAppName(), instanceInfo.getVIPAddress(),
                instanceInfo.getIPAddr(), instanceInfo.getPort(), isReplication);

        int leaseDuration = Lease.DEFAULT_DURATION_IN_SECS;
        if (instanceInfo.getLeaseInfo() != null && instanceInfo.getLeaseInfo().getDurationInSecs() > 0) {
            leaseDuration = instanceInfo.getLeaseInfo().getDurationInSecs();
        }
        publishEvent(new EurekaInstanceRegisteredEvent(this, instanceInfo, leaseDuration, isReplication));

        // 如果IP白名单为空，则表示不进行限制
        // IP白名单不为空时，只允许指定 IP 进行注册.
        if (CollectionUtils.isEmpty(allowedRegisteredIpAddress) || allowedRegisteredIpAddress.contains(instanceInfo.getIPAddr())) {
            log.info("\n IP Address Allows Registration: {}", instanceInfo.getIPAddr());
            super.register(instanceInfo, isReplication);
            return;
        }
        log.warn("\n IP Address Not Allows Registration: {}", instanceInfo.getIPAddr());
    }

    /**
     * 取消实例注册
     * 当客户机关闭通知服务器从流量中删除实例时,通常会调用它。
     *
     * @param appName       ：应用名称
     * @param serverId      ：服务Id，实例的唯一标识符
     * @param isReplication ：如果这是来自其他节点的复制事件,则是 false。如果实例被删除，则为 true.
     * @return
     */
    @Override
    public boolean cancel(String appName, String serverId, boolean isReplication) {
        log.warn("cancel: appName={}, serverId={}, isReplication={}.", appName, serverId, isReplication);
        publishEvent(new EurekaInstanceCanceledEvent(this, appName, serverId, isReplication));
        return super.cancel(appName, serverId, isReplication);
    }

    /**
     * 将给定的应用程序名称标记为更新,并标记它是否起源于复制。
     *
     * @param appName       ：应用名称
     * @param serverId      ：服务Id，实例的唯一标识符
     * @param isReplication ：是否复制
     * @return
     */
    @Override
    public boolean renew(final String appName, final String serverId, boolean isReplication) {
        log.debug("renew: appName={}, serverId={}, isReplication={}.", appName, serverId, isReplication);
        List<Application> applications = getSortedApplications();
        for (Application input : applications) {
            if (input.getName().equals(appName)) {
                InstanceInfo instance = null;
                for (InstanceInfo info : input.getInstances()) {
                    if (info.getId().equals(serverId)) {
                        instance = info;
                        break;
                    }
                }
                publishEvent(new EurekaInstanceRenewedEvent(this, appName, serverId, instance, isReplication));
                break;
            }
        }
        return super.renew(appName, serverId, isReplication);
    }

    /**
     * 内部取消
     *
     * @param appName       ：服务名称
     * @param id            ：服务ID
     * @param isReplication ：是否复制
     * @return
     */
    @Override
    protected boolean internalCancel(String appName, String id, boolean isReplication) {
        log.warn("internalCancel: appName={}, serverId={}, isReplication={}.", appName, id, isReplication);
        publishEvent(new EurekaInstanceCanceledEvent(this, appName, id, isReplication));
        return super.internalCancel(appName, id, isReplication);
    }

    private void publishEvent(ApplicationEvent applicationEvent) {
        this.applicationContext.publishEvent(applicationEvent);
    }

}

